import json
import os
import uuid
from typing import Optional

from data_preprocess.doc_split import slice_text_for_rag
from data_preprocess.pdf import extract_text_from_pdf
from data_preprocess.txt import extract_text_from_txt
from data_preprocess.word import extract_text_from_doc, extract_text_from_docx
from util.common import remove_markdown_code_block
from util.es import delete_by_id, insert_data
from util.llm import http_request

REMOVE_SYSTEM_PROMPT = """任务：
以下是一段文本，是使用程序从pdf/word中提取的，可能存在多出换行符、带有乱码等情况。请帮我将多余字符去掉。
输出要求：
切记提取后的文本必须和原文保持一致。只需输出文章正文，不要添加任何解释性文字。
正文：
{}
"""


class File:
    def __init__(self, file_id: str, file_name: str, permission_type: int, permission_set: list[str], file: bytes,
                 file_path=None):
        self.file_id = file_id
        self.file_name = file_name
        self.permission_set = permission_set
        self.permission_type = permission_type  # 权限类型，为1则所有人均可访问，为2时只有permission_set中的人可以访问
        self.file = file
        self.content = ""
        self.trunks = []
        self.inner_id = str(uuid.uuid4())
        self.file_path = file_path
        self.file_type = self.detect_file_type()

    async def remove_garbled_characters(self, content):
        if len(content) == 0:
            return ""
        system_prompt = REMOVE_SYSTEM_PROMPT.format(content)
        # r = await http_request("你是一个语言专家", [], system_prompt, model="hunyuan-turbo")
        r = await http_request("你是一个语言专家", [], system_prompt, model="deepseek-chat")
        return r

    async def upload_to_es(self):
        print(f"🔍 upload_to_es开始，文本块数: {len(self.trunks) if self.trunks else 0}")
        
        if not self.trunks:
            print(f"❌ 没有文本块可以上传")
            return
        
        data = [{"title": self.file_name, "content": x["text"], "file_id": self.file_id,
                 "permission_set": self.permission_set} for x in self.trunks]
        
        print(f"🔍 准备上传 {len(data)} 个数据块")
        
        for i, x in enumerate(data):
            x["content"] = remove_markdown_code_block(x["content"])
            print(f"🔍 数据块 {i+1}: {x['content'][:100]}...")
        
        print(f"🔍 调用insert_data上传到ES...")
        insert_data("law1", data)
        print(f"✅ upload_to_es完成")
        # insert_data("law", [{"title": "中国联通合作方数据安全管理规范", "content": x} for x in law_data])

    def delete_from_es(self):
        result, info = delete_by_id("law1", self.file_id)
        if not result:
            print("删除失败：{}".format(info))

    def upload_to_vdb(self):
        pass

    def delete_from_vdb(self):
        pass

    def process(self):
        pass

    def detect_file_type(self) -> Optional[str]:
        """检测文件类型（基于文件头）"""
        if self.file_path:
            # 从文件路径检测
            with open(self.file_path, 'rb') as f:
                header = f.read(1024)
        else:
            # 从字节数据检测
            header = self.file[:1024] if len(self.file) >= 1024 else self.file
            
        # 简单的文件头检测逻辑
        if b'%PDF-' in header:
            return 'pdf'
        elif header.startswith(b'\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1'):
            # return 'doc'
            return None
        elif header.startswith(b'PK\x03\x04'):
            return 'docx'
        
        # 尝试解码为文本
        try:
            if self.file_path:
                with open(self.file_path, 'r', encoding='utf-8') as f:
                    f.read(1024)  # 尝试读取前1024字节
            else:
                # 尝试解码字节数据
                self.file[:1024].decode('utf-8')
            return "txt"
        except UnicodeDecodeError:
            # 如果UTF-8解码失败，尝试其他编码
            try:
                if self.file_path:
                    with open(self.file_path, 'r', encoding='gbk') as f:
                        f.read(1024)
                else:
                    self.file[:1024].decode('gbk')
                return "txt"
            except UnicodeDecodeError:
                return None

    async def file_extract(self, file_type, use_simple_extract=False):
        if self.file_path:
            # 从文件路径提取
            if file_type == "pdf":
                if use_simple_extract:
                    from data_preprocess.pdf_simple import extract_text_from_pdf_simple_async
                    content = await extract_text_from_pdf_simple_async(self.file_path)
                else:
                    content = await extract_text_from_pdf(self.file_path)
            elif file_type == "docx":
                content = extract_text_from_docx(self.file_path)
            elif file_type == "doc":
                content = extract_text_from_doc(self.file_path)
            else:
                content = extract_text_from_txt(self.file_path)
        else:
            # 从字节数据提取
            import tempfile
            with tempfile.NamedTemporaryFile(delete=False, suffix=f'.{file_type}') as temp_file:
                temp_file.write(self.file)
                temp_file_path = temp_file.name
            
            try:
                if file_type == "pdf":
                    if use_simple_extract:
                        from data_preprocess.pdf_simple import extract_text_from_pdf_simple_async
                        content = await extract_text_from_pdf_simple_async(temp_file_path)
                    else:
                        content = await extract_text_from_pdf(temp_file_path)
                elif file_type == "docx":
                    content = extract_text_from_docx(temp_file_path)
                elif file_type == "doc":
                    content = extract_text_from_doc(temp_file_path)
                else:
                    content = extract_text_from_txt(temp_file_path)
            finally:
                # 清理临时文件
                try:
                    os.unlink(temp_file_path)
                except:
                    pass

        print(content)
        self.content = content if content is not None and content != "" else ""

    def content_trunk(self, chunk_size=1000, overlap=0):
        self.trunks = slice_text_for_rag(self.content, chunk_size, overlap)

    def file_save(self):
        # save_dir = "uploads"
        # os.makedirs(save_dir, exist_ok=True)
        # self.save_path = os.path.join(save_dir, self.inner_id)
        # with open(self.save_path, "wb") as f:
        #     f.write(self.file)
        pass
